#include "GGXMaterial.cuh"

using namespace OBR;

#define GGX_USE_IS

/**
 * @brief
 *
 * @param alpha ggx param
 * @param m microfacet normal
 * @param n macrofacet normal
 */
static __forceinline__ __device__ float ggx_D(float alpha, float3 m, float3 n)
{
    float sqr_alpha = alpha * alpha;
    float tan2_theta = tan2(m, n);
    float cos4_theta = cos2(m, n) * cos2(m, n);
    float e = (1.0f / sqr_alpha) * tan2_theta;
    return (1.0f / (M_PIf * sqr_alpha * cos4_theta * (1 + e) * (1 + e)));
}

/**
 * @brief
 *
 * @param alpha ggx param
 * @param tan_theta tangent value of angle between microfacet normal & light path
 */
static __forceinline__ __device__ float ggx_lambda(float alpha, float tan_theta)
{
    tan_theta = fmaxf(tan_theta, 0.0f);
    return 0.5 * (sqrtf(alpha * alpha * tan_theta * tan_theta + 1) - 1);
}

/**
 * @brief
 *
 * @param alpha ggx param
 * @param i in-going direction
 * @param o out-going direction
 * @param n macrofacet normal
 */
static __forceinline__ __device__ float ggx_G(float alpha, float3 i, float3 o, float3 n)
{
    float tan_theta_i, tan_theta_o;
    tan_theta_i = tan(i, n);
    tan_theta_o = tan(o, n);
    return 1.0f / (1.0f + ggx_lambda(alpha, tan_theta_i) + ggx_lambda(alpha, tan_theta_o));
}

/**
 * @brief
 *
 * @param alpha ggx param
 * @param i in-going direction
 * @param o out-going direction
 * @param n macrofacet normal
 */
static __forceinline__ __device__ float ggx_pdf(float alpha, float3 o, float3 i, float3 n)
{
    float3 m = normalize(i + o);
    float D = ggx_D(alpha, m, n);
    float G1 = 1.0f / (1.0f + ggx_lambda(alpha, tan(o, n)));

    /* visible */
    // return D * G1 * fabs(dot(o, m)) / fabs(dot(o, n)) / (4.0f * dot(o, m));
    return 0.25f * D * G1 / fabs(dot(o, n));

    // /* all */
    // return D * fabs(dot(o, n)) / (4.0f * dot(o, m));
}

/**
 * @brief
 *
 * @param alpha ggx param
 * @param i in-going direction
 * @param o out-going direction
 * @param n macrofacet normal
 */
static __forceinline__ __device__ float3 ggx_sample_wh(float alpha, float3 i, float3 n, float u1,
                                                       float u2)
{
    Onb onb(n);
    float3 wi_local = onb.world2local(i);
    float3 wi_stretched = normalize(wi_local);

    /* ---------------- simulate P22_{wi}(x_slope, y_slope, 1, 1) --------------- */
    float cos_theta = cos(i, n);
    float sin_theta = sin(i, n);
    float tan_theta = tan(i, n);
    float slope_x, slope_y;
    if (cos_theta > .9999)
    {
        // near perfect sample
        float r = sqrtf(u1 / 1.0f - u1);
        float phi = 2.0f * M_PIf * u2;
        slope_x = r * cos(phi);
        slope_y = r * sin(phi);
    }
    else
    {
        float a = 1 / tan_theta;
        float G1 = 1.0f / (1.0f + ggx_lambda(1.0, tan_theta));
        // sample slope_x
        float A = 2 * u1 / G1 - 1;
        float tmp = 1.f / (A * A - 1.f);
        if (tmp > 1e10)
            tmp = 1e10;
        float B = tan_theta;
        float D = sqrtf(fmaxf(float(B * B * tmp * tmp - (A * A - B * B) * tmp), 0.0f));
        float slope_x_1 = B * tmp - D;
        float slope_x_2 = B * tmp + D;
        slope_x = (A < 0 || slope_x_2 > 1.f / tan_theta) ? slope_x_1 : slope_x_2;
        // sample slope_y
        float S;
        if (u2 > 0.5f)
        {
            S = 1.f;
            u2 = 2.f * (u2 - .5f);
        }
        else
        {
            S = -1.f;
            u2 = 2.f * (.5f - u2);
        }
        float z = (u2 * (u2 * (u2 * 0.27385f - 0.73369f) + 0.46341f)) /
                  (u2 * (u2 * (u2 * 0.093073f + 0.309420f) - 1.000000f) + 0.597999f);
        slope_y = S * z * sqrtf(1.f + slope_x * slope_x);
    }

    /* --------------------------------- rotate --------------------------------- */
    float tmp = local_cos_phi(wi_stretched) * slope_x - local_sin_phi(wi_stretched) * slope_y;
    slope_y = local_sin_phi(wi_stretched) * slope_x + local_cos_phi(wi_stretched) * slope_y;
    slope_x = tmp;

    /* ----------------------------- transform back ----------------------------- */
    slope_x = alpha * slope_x;
    slope_y = alpha * slope_y;
    float3 wh_local = normalize(make_float3(-slope_x, -slope_y, 1.0f));
    return onb.local2world(wh_local);
}

/**
 * @brief calc combined effect of bsdf / cos theta
 *
 */
static __forceinline__ __device__ float calcPdf(SamplePRD *prd, GGXMaterialData *data)
{
#ifdef GGX_USE_IS
    InteractionData &si = prd->inter;
    float rad = si.has_differentials ? sqrtf(min(si.dudx * si.dudx + si.dudy * si.dudy,
                                                 si.dvdx * si.dvdx + si.dvdy * si.dvdy))
                                     : 0;
    float alpha = data->alpha.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE);
    return ggx_pdf(alpha, -si.wo.d, si.wi.d, si.n);
#else
    return M_1_PIf * max(0.0f, dot(prd->inter.n, prd->inter.wi.d));
#endif
}

/**
 * @brief calc combined effect of bsdf / cos theta
 *
 */
static __forceinline__ __device__ Spectrum calcFr(SamplePRD *prd, float *pdf, GGXMaterialData *data)
{
    InteractionData &si = prd->inter;

    // differential
    float rad = si.has_differentials ? sqrtf(min(si.dudx * si.dudx + si.dudy * si.dudy,
                                                 si.dvdx * si.dvdx + si.dvdy * si.dvdy))
                                     : 0;

    // shade
    Spectrum beta(1.0f);
    float3 i = si.wi.d;
    float3 o = -si.wo.d;
    float3 n = si.n;
    float3 m = normalize(i + o);
    float alpha = data->alpha.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE);
    Spectrum albedo = data->albedo.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE);

    float bsdf = 1.0f;
    bsdf *= fr_dielectric(clamp(dot(o, m), 0.0f, 1.0f), data->eta);
    bsdf *= ggx_G(alpha, o, i, n);
    bsdf *= ggx_D(alpha, m, n);
    bsdf /= 4.0f * fabs(dot(i, n)) * fabs(dot(o, n));
    beta *= albedo;
    beta *= bsdf;
    beta *= max(0.0f, dot(i, n));

    *pdf = ggx_pdf(alpha, o, i, n);
    return beta.clamp();
}

/**
 * @brief generate new ray info and add to the next ray, then return beta of the ray
 *
 * @param prd
 * @param instance_id
 * @return beta
 */
static __forceinline__ __device__ Spectrum rayGeneration(SamplePRD *prd, float *pdf,
                                                         GGXMaterialData *data, float u0, float u1)
{
    InteractionData &si = prd->inter;
    float rad = si.has_differentials ? sqrtf(min(si.dudx * si.dudx + si.dudy * si.dudy,
                                                 si.dvdx * si.dvdx + si.dvdy * si.dvdy))
                                     : 0;
    float alpha = data->alpha.sample(si.uv, rad, rad, MipmapSampleMode::SQUARE);

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("alpha: %f\n", alpha);
    }
#endif

#ifdef GGX_USE_IS
    float3 m = ggx_sample_wh(alpha, -si.wo.d, si.n, u0, u1);
    si.wi = si.spawnRay(reflect(si.wo.d, m));
#else
    // sample hemisphere
    Onb onb(prd->inter.n);
    float3 hemisphere_dir = cosine_sample_hemisphere(u0, u1);
    prd->inter.wi = prd->inter.spawnRay(onb.local2world(hemisphere_dir));
#endif

    return calcFr(prd, pdf, data);
}

extern "C" __device__ Spectrum __direct_callable__ggx_material(SamplePRD *prd, float *pdf, float u0,
                                                               float u1, MaterialCallMode mode)
{
    // get sbt data
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    GGXMaterialData *data = reinterpret_cast<GGXMaterialData *>(sbt_record->callable_payload);
    data = data + prd->inter.mat_instance_id;

    if (mode == MaterialCallMode::SAMPLE_FR)
    {
        return rayGeneration(prd, pdf, data, u0, u1);
    }
    else if (mode == MaterialCallMode::CALC_FR_PDF)
    {
        return calcFr(prd, pdf, data);
    }
    return Spectrum(0.0f);
}